#########################################################################################
##  Stage 2: build DALI wheels on top of the dependencies image built in Stage 1
#########################################################################################
ARG DEPS_IMAGE_NAME
# clean builder without source code inside
FROM ${DEPS_IMAGE_NAME} as builder

ARG PYVER=3.10
ARG PYV=310

ENV PYVER=${PYVER} \
    PYV=${PYV} \
    PYTHONPATH=/opt/python/v \
    PYBIN=${PYTHONPATH}/bin \
    PYLIB=${PYTHONPATH}/lib \
    PATH=/opt/python/cp${PYV}-cp${PYV}/bin:/opt/python/cp39-cp39/bin:${PYBIN}:/opt/python/cp310-cp310/bin:/opt/python/cp311-cp311/bin:${PYBIN}:/opt/python/cp312-cp312/bin:/opt/python/cp313-cp313/bin:${PYBIN}:${PATH} \
    LD_LIBRARY_PATH=/opt/python/cp${PYV}-cp${PYV}/lib:/usr/local/lib:/opt/dali/${DALI_BUILD_DIR}:/opt/python/cp39-cp39/lib:/opt/python/cp310-cp310/lib:/opt/python/cp311-cp311/lib:/opt/python/cp312-cp312/lib:/opt/python/cp313-cp313/lib:${PYLIB}:${LD_LIBRARY_PATH} \
    LIBRARY_PATH=/opt/python/cp${PYV}-cp${PYV}/lib:/usr/local/lib:/opt/dali/${DALI_BUILD_DIR}:/opt/python/cp39-cp39/lib:/opt/python/cp310-cp310/lib:/opt/python/cp311-cp311/lib:/opt/python/cp312-cp312/lib:/opt/python/cp313-cp313/lib:${PYLIB}:${LIBRARY_PATH}

RUN ln -s /opt/python/cp${PYV}* /opt/python/v

# install Python bindings and patch it to use the clang we have here
RUN pip install future setuptools wheel clang==14.0 flake8 bandit astunparse gast dm-tree "black[jupyter]"==25.11.0 nvtx makefun && \
    PY_CLANG_PATH=$(echo $(pip show clang) | sed 's/.*Location: \(.*\) Requires.*/\1/')/clang/cindex.py && \
    LIBCLANG_PATH=/usr/lib64/libclang.so && \
    sed -i "s|library_file = None|library_file = \"${LIBCLANG_PATH}\"|" ${PY_CLANG_PATH} && \
    rm -rf /root/.cache/pip/

RUN ldconfig

WORKDIR /opt/dali

ARG CC
ARG CXX
ENV CC=${CC} \
    CXX=${CXX}
# Optional build arguments

ARG ARCH
ENV ARCH=${ARCH}
ARG CUDA_TARGET_ARCHS
ENV CUDA_TARGET_ARCHS=${CUDA_TARGET_ARCHS}
ARG CMAKE_BUILD_TYPE
ENV CMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
ARG BUILD_TEST
ENV BUILD_TEST=${BUILD_TEST}
ARG BUILD_BENCHMARK
ENV BUILD_BENCHMARK=${BUILD_BENCHMARK}
ARG BUILD_NVTX
ENV BUILD_NVTX=${BUILD_NVTX}
ARG BUILD_PYTHON
ENV BUILD_PYTHON=${BUILD_PYTHON}
ARG BUILD_LMDB
ENV BUILD_LMDB=${BUILD_LMDB}
ARG BUILD_JPEG_TURBO
ENV BUILD_JPEG_TURBO=${BUILD_JPEG_TURBO}
ARG BUILD_NVJPEG
ENV BUILD_NVJPEG=${BUILD_NVJPEG}
ARG BUILD_NVJPEG2K
ENV BUILD_NVJPEG2K=${BUILD_NVJPEG2K}
ARG BUILD_LIBTIFF
ENV BUILD_LIBTIFF=${BUILD_LIBTIFF}
ARG BUILD_LIBSND
ENV BUILD_LIBSND=${BUILD_LIBSND}
ARG BUILD_LIBTAR
ENV BUILD_LIBTAR=${BUILD_LIBTAR}
ARG BUILD_FFTS
ENV BUILD_FFTS=${BUILD_FFTS}
ARG BUILD_CFITSIO
ENV BUILD_CFITSIO=${BUILD_CFITSIO}
ARG BUILD_NVOF
ENV BUILD_NVOF=${BUILD_NVOF}
ARG BUILD_NVDEC
ENV BUILD_NVDEC=${BUILD_NVDEC}
ARG BUILD_NVML
ENV BUILD_NVML=${BUILD_NVML}
ARG BUILD_CUFILE
ENV BUILD_CUFILE=${BUILD_CUFILE}
ARG BUILD_NVCOMP
ENV BUILD_NVCOMP=${BUILD_NVCOMP}
ARG BUILD_CVCUDA
ENV BUILD_CVCUDA=${BUILD_CVCUDA}
ARG BUILD_NVIMAGECODEC
ENV BUILD_NVIMAGECODEC=${BUILD_NVIMAGECODEC}
ARG LINK_DRIVER
ENV LINK_DRIVER=${LINK_DRIVER}
ARG WITH_DYNAMIC_CUDA_TOOLKIT
ENV WITH_DYNAMIC_CUDA_TOOLKIT=${WITH_DYNAMIC_CUDA_TOOLKIT}
ARG WITH_DYNAMIC_NVJPEG
ENV WITH_DYNAMIC_NVJPEG=${WITH_DYNAMIC_NVJPEG:-ON}
ARG WITH_DYNAMIC_CUFFT
ENV WITH_DYNAMIC_CUFFT=${WITH_DYNAMIC_CUFFT:-ON}
ARG WITH_DYNAMIC_NPP
ENV WITH_DYNAMIC_NPP=${WITH_DYNAMIC_NPP:-ON}
ARG WITH_DYNAMIC_NVIMGCODEC
ENV WITH_DYNAMIC_NVIMGCODEC=${WITH_DYNAMIC_NVIMGCODEC:-ON}
ARG WITH_DYNAMIC_NVCOMP
ENV WITH_DYNAMIC_NVCOMP=${WITH_DYNAMIC_NVCOMP:-ON}
ARG BUILD_WITH_ASAN
ENV BUILD_WITH_ASAN=${BUILD_WITH_ASAN}
ENV BUILD_WITH_LSAN=${BUILD_WITH_LSAN}
ARG BUILD_WITH_UBSAN
ENV BUILD_WITH_UBSAN=${BUILD_WITH_UBSAN}
ARG PYTHON_VERSIONS
ENV PYTHON_VERSIONS=${PYTHON_VERSIONS}
ARG STRIP_BINARY
ENV STRIP_BINARY=${STRIP_BINARY}
ARG VERBOSE_LOGS
ENV VERBOSE_LOGS=${VERBOSE_LOGS}
ARG NVIDIA_DALI_BUILD_FLAVOR
ENV NVIDIA_DALI_BUILD_FLAVOR=${NVIDIA_DALI_BUILD_FLAVOR}
ARG GIT_SHA
ENV GIT_SHA=${GIT_SHA}
ARG DALI_TIMESTAMP
ENV DALI_TIMESTAMP=${DALI_TIMESTAMP}
ARG WHL_PLATFORM_NAME
ENV WHL_PLATFORM_NAME=${WHL_PLATFORM_NAME}
ARG WHL_COMPRESSION
ENV WHL_COMPRESSION=${WHL_COMPRESSION}
ARG EXTRA_CMAKE_OPTIONS
ENV EXTRA_CMAKE_OPTIONS=${EXTRA_CMAKE_OPTIONS}

ARG NVIDIA_BUILD_ID
ENV NVIDIA_BUILD_ID ${NVIDIA_BUILD_ID:-0}
RUN mkdir /wheelhouse && chmod 0777 /wheelhouse

FROM builder
COPY . .

ARG DALI_BUILD_DIR=build-docker-release
WORKDIR /opt/dali/${DALI_BUILD_DIR}

RUN bash /opt/dali/docker/build_helper.sh
